import os
import time
from pymilvus import MilvusClient, model as milvus_model
from openai import OpenAI

# 1. 加载知识库
def load_knowledge(file_path):
    import re
    with open(file_path, encoding='utf-8') as f:
        text = f.read()
    # 按“**第xxx条**”分割
    docs = re.split(r"\*\*第[一二三四五六七八九零百千\d]+条\*\*", text)
    docs = [d.strip() for d in docs if d.strip()]
    return docs

docs = load_knowledge("mfd.md")

# 2. 构建嵌入模型
embedding_model = milvus_model.DefaultEmbeddingFunction()
doc_embeddings = embedding_model.encode_documents(docs)

# 3. 连接Milvus Lite并创建collection
collection_name = "property_law_rag"
embedding_dim = len(doc_embeddings[0])
milvus_client = MilvusClient(uri="./milvus_property_law.db")

# 如果collection已存在则删除
if milvus_client.has_collection(collection_name):
    milvus_client.drop_collection(collection_name)

milvus_client.create_collection(
    collection_name=collection_name,
    dimension=embedding_dim,
    metric_type="IP",
    consistency_level="Strong"
)

# 4. 插入数据
data = []
for i, doc in enumerate(docs):
    data.append({"id": i, "vector": doc_embeddings[i], "text": doc})
milvus_client.insert(collection_name=collection_name, data=data)

# 5. 定义RAG检索+生成函数
def rag_ask(question, top_k=3):
    # 检索
    q_emb = embedding_model.encode_queries([question])
    search_res = milvus_client.search(
        collection_name=collection_name,
        data=q_emb,
        limit=top_k,
        search_params={"metric_type": "IP", "params": {}},
        output_fields=["text"],
    )
    context = "\n".join([res["entity"]["text"] for res in search_res[0]])
    # 组装prompt
    prompt = f"""你是中国民法典物权篇的智能助手，请结合以下内容回答用户问题。\n知识片段：\n{context}\n\n问题：{question}\n\n请用中文简明回答。"""
    # 调用deepseek-chat
    api_key = os.getenv("DEEPSEEK_API_KEY")
    client = OpenAI(api_key=api_key, base_url="https://api.deepseek.com/v1")
    response = client.chat.completions.create(
        model="deepseek-chat",
        messages=[
            {"role": "system", "content": "你是中国民法典物权篇的智能助手。"},
            {"role": "user", "content": prompt}
        ]
    )
    return response.choices[0].message.content

# 6. 测试两个问题并计时
if __name__ == "__main__":
    questions = [
        "不动产物权变更需要登记吗？",
        "拾得遗失物应当如何处理？"
    ]
    for q in questions:
        print(f"\n问题：{q}")
        start = time.time()
        answer = rag_ask(q)
        end = time.time()
        print("答案：", answer)
        print("耗时：{:.2f}秒".format(end - start))